#include "Process.h"
#include "RemoteThread.h"
#include <TlHelp32.h>
#include <stdexcept>

#define PROCESS_ACCESS_FLAGS PROCESS_CREATE_THREAD | PROCESS_VM_OPERATION | SYNCHRONIZE | PROCESS_VM_READ | PROCESS_VM_WRITE | PROCESS_QUERY_INFORMATION   // all the flags needed for the stuff we gonna do (R/W memory, create threads..)
#define LOAD_LIBRARY_A_STR "LoadLibraryA"

Process::Process(DWORD processID):
m_processHandle(NULL)
{
	HANDLE processHandle = OpenProcess(PROCESS_ACCESS_FLAGS, FALSE, processID);
	if (processHandle == NULL)
		throw std::runtime_error("Process::Process - Process could not be opened.");
	m_processHandle = processHandle;
}
Process::Process(const char* processName):
m_processHandle(NULL)
{
	DWORD processID = GetProcessIdByName(processName);
	if (processID == 0)
		throw std::runtime_error("Process::Process - Process with given name not found.");

	HANDLE processHandle = OpenProcess(PROCESS_ACCESS_FLAGS, FALSE, processID);
	if (processHandle == NULL)
		throw std::runtime_error("Process::Process - Process could not be opened.");
	m_processHandle = processHandle;
}
Process::Process(Process&& rhs):
m_processHandle(rhs.m_processHandle),
m_allocs(std::move(rhs.m_allocs))
{
	rhs.m_processHandle = NULL;
}
Process::~Process(void)
{
	m_allocs.clear();
	if (m_processHandle != NULL)
		CloseHandle(m_processHandle);
}
Process& Process::operator=(Process&& rhs)
{
	m_processHandle = rhs.m_processHandle;
	m_allocs = std::move(rhs.m_allocs);

	rhs.m_processHandle = NULL;

	return *this;
}


/*
Finds the >> FIRST << process with given name and returns its ID
*/
DWORD Process::GetProcessIdByName(const char* processName)
{
	HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (snapshot == INVALID_HANDLE_VALUE)
		throw std::runtime_error("Process:GetProcessIdFromName - Creation of processes snapshot failed.");

	PROCESSENTRY32 processEntry = { 0 };
	processEntry.dwSize = sizeof(PROCESSENTRY32);
	DWORD processID = 0;

	bool processFound = Process32First(snapshot, &processEntry);

	while (processFound)
	{
		if (strcmp(processName, processEntry.szExeFile) == 0)
		{
			processID = processEntry.th32ProcessID;
			break;
		}
		processFound = Process32Next(snapshot, &processEntry);
	}
	CloseHandle(snapshot);
	return processID;
}

/*
Simple VirtualAllocEx wrapper with additional error checking
*/
MemoryAlloc& Process::AllocMemory(DWORD size)
{
	MemoryAlloc alloc(m_processHandle, size);
	m_allocs.push_back(std::move(alloc));
	return m_allocs.back();
}

/*
Simple CreateRemoteThread wrapper with additional error checking
*/
DWORD Process::CallFunction(DWORD(__stdcall*fn)(void*), void* arg)
{
	RemoteThread thread(m_processHandle, fn);
	return thread.Call(arg);
}

/*
Most basic library injection method using the Windows LoadLibraryA API. DllMain is called automatically, if the library is .dll if its an .exe, nothing is called
Very reliable, but also leaves all the evidence in the target process - (Library is added to the PEB LDR data linked list, the memory page names are same as the dll name etc.)
*/
void Process::InjectLibrary(const char* libraryPath)
{
	DWORD fileAttributes = GetFileAttributesA(libraryPath);
	if (fileAttributes == INVALID_FILE_ATTRIBUTES)
		throw std::runtime_error("Process::InjectLibrary - Library file not found.");

	DWORD pathSize = strlen(libraryPath) + 1;
	MemoryAlloc libPathAlloc(m_processHandle, pathSize);
	libPathAlloc.Write(0, libraryPath, pathSize);
	FARPROC loadLibraryA = GetProcAddress(GetModuleHandleA("kernel32"), "LoadLibraryA");
	CallFunction(reinterpret_cast<DWORD(__stdcall*)(void*)>(loadLibraryA), libPathAlloc.GetAddress());
}

/*
Abit more advanced method, that injects the image manually. The PE headers and all sections are copied to the target process using WriteProcessMemory
This method should be more stealthy. Not sure about reliability, but should work most of the time
Relocations should be present in the image, since we dont know at what address our image is gonna end up
*/
void Process::InjectModuleRaw(void* moduleBase, DWORD(__stdcall* functionAddress)(void*))
{
	IMAGE_DOS_HEADER* dosHeader = static_cast<IMAGE_DOS_HEADER*>(moduleBase);
	IMAGE_NT_HEADERS* ntHeaders = reinterpret_cast<IMAGE_NT_HEADERS*>(static_cast<char*>(moduleBase) + dosHeader->e_lfanew);
	IMAGE_SECTION_HEADER* sectionHeader = reinterpret_cast<IMAGE_SECTION_HEADER*>(reinterpret_cast<char*>(&ntHeaders->OptionalHeader) + ntHeaders->FileHeader.SizeOfOptionalHeader);
	
	MemoryAlloc alloc(m_processHandle, ntHeaders->OptionalHeader.SizeOfImage); //Alloc memory for the whole image
	alloc.Write(0, moduleBase, ntHeaders->OptionalHeader.SizeOfHeaders); // Copy the PE headers
	
	
	for (int i = 0; i < ntHeaders->FileHeader.NumberOfSections; i++) // Copy the sections one by one
	{
		DWORD virtualAddress = sectionHeader[i].VirtualAddress;
		alloc.Write(virtualAddress, reinterpret_cast<char*>(moduleBase) + virtualAddress, sectionHeader[i].Misc.VirtualSize); 
	}

	IMAGE_BASE_RELOCATION* baseReloc = 0; // Find the reloc directory and if exists, fix up the relocations
	baseReloc = reinterpret_cast<IMAGE_BASE_RELOCATION*>(reinterpret_cast<char*>(moduleBase)+ ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress); 
	if (baseReloc != 0)
	{
		while (baseReloc->VirtualAddress != 0)
		{
			WORD* reloc = reinterpret_cast<WORD*>(baseReloc + 1);
			for (int i = 0; i * 2 < baseReloc->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION); i++)
			{
				if (reloc[i] >> 12 == IMAGE_REL_BASED_HIGHLOW)
				{
					int baseDelta = reinterpret_cast<DWORD>(alloc.GetAddress()) - reinterpret_cast<DWORD>(moduleBase);
					DWORD relocOffset = (reloc[i] & 0x0FFF) + baseReloc->VirtualAddress;
					DWORD val = 0;
					alloc.Read(relocOffset, &val, 4);
					DWORD newVal = val + baseDelta;
					alloc.Write(relocOffset, reinterpret_cast<const void*>(&newVal), 4);
				}
			}
			baseReloc = reinterpret_cast<IMAGE_BASE_RELOCATION*>(reinterpret_cast<char*>(baseReloc) + baseReloc->SizeOfBlock);
		}
	}


	// This doesnt need to be there, but i felt like this might be better than simply calling the actual entry point
	// Calling nothing at all and let the user do it later on could be possibility aswell, but w/e gonna keep it like this for now :)

	DWORD entryPointOffset = reinterpret_cast<DWORD>(functionAddress) - reinterpret_cast<DWORD>(moduleBase);
	DWORD(__stdcall* entryPointFunc)(void*) = reinterpret_cast<DWORD(__stdcall*)(void*)>(reinterpret_cast<char*>(alloc.GetAddress()) + entryPointOffset);
	CallFunction(entryPointFunc, 0);
}